#include "Basis.h"
#include "Notification.h"

LARGE_INTEGER g_CallbackCookie;

VOID InitializeNotification()
{
	// register process notify routine
	NTSTATUS status = PsSetCreateProcessNotifyRoutineEx(ProcessNotifyRoutineEx, FALSE); // required /IntegrityCheck in link.exe
	if (!NT_SUCCESS(status))
	{
		KdPrint(("failed to register process notify routine, status 0x%x", status));
	}

	// register image notify routine
	status = PsSetLoadImageNotifyRoutine(LoadImageNotifyRoutine);
	if (!NT_SUCCESS(status))
	{
		KdPrint(("failed to register image notify routine, status 0x%x", status));
	}

	// register register call back
	status = CmRegisterCallback(RegistryCallback, NULL, &g_CallbackCookie);
	if (!NT_SUCCESS(status))
	{
		KdPrint(("failed to register register call back, status 0x%x", status));
	}
}

VOID DisposeNotification()
{
	// remove image notify routine
	PsRemoveLoadImageNotifyRoutine(LoadImageNotifyRoutine);

	// remove process notify routine
	PsSetCreateProcessNotifyRoutineEx(ProcessNotifyRoutineEx, TRUE);

	// remove register call back
	CmUnRegisterCallback(g_CallbackCookie);
}

VOID ProcessNotifyRoutineEx(_Inout_ PEPROCESS Process, _In_ HANDLE ProcessId, _Inout_opt_ PPS_CREATE_NOTIFY_INFO CreateInfo)
{
	UNREFERENCED_PARAMETER(Process);

	if (CreateInfo != NULL)
	{
		KdPrint(("Creating ProcessId %d, Status %d, File %wZ, Command %wZ", ProcessId, CreateInfo->CreationStatus, CreateInfo->ImageFileName, CreateInfo->CommandLine));
	}
	else
	{
		KdPrint(("Closing ProcessId %d", ProcessId));
	}
}

VOID LoadImageNotifyRoutine(_In_ PUNICODE_STRING FullImageName, _In_ HANDLE ProcessId, _In_ PIMAGE_INFO ImageInfo)
{
	KdPrint(("ProcessId %d, ImageBase 0x%x, Name %wZ", ProcessId, ImageInfo->ImageBase, FullImageName));
}

NTSTATUS RegistryCallback(IN PVOID CallbackContext, IN PVOID Argument1, IN PVOID Argument2)
{
	UNREFERENCED_PARAMETER(CallbackContext);
	UNREFERENCED_PARAMETER(Argument2);

	REG_NOTIFY_CLASS type = (REG_NOTIFY_CLASS)Argument1;
	//KdPrint(("Reister call back type %d", type));
	UNREFERENCED_PARAMETER(type);

	return STATUS_SUCCESS;
}